import torch
from torch.utils.data import DataLoader
from tqdm import tqdm
from colorama import Fore
import pickle
class Eval:
    def __init__(self,model,test_dataset,collate_fn):
        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
        kwargs = {'num_workers':1,'pin_memory':True} if torch.cuda.is_available() else {}

        self.model = model.to(self.device)
        self.model.eval()
        self.test_dataloader = DataLoader(
                                test_dataset,
                                1,
                                shuffle=False,
                                collate_fn = collate_fn,
                                **kwargs
                                )
        
        self.image_features = {}
        self.image_person_ids = {}
    

    @torch.no_grad()
    def run(self):
        pbar = tqdm(self.test_dataloader,
                    desc=f'Test:',
                    bar_format='{l_bar}%s{bar}%s{r_bar}' % (Fore.BLUE, Fore.RESET)
                    )
        for images,person_ids,camera_ids,video_ids,image_paths in pbar:
            images = images.to(self.device)
            features = self.model(images).cpu().numpy()
            for feature,image_path,person_id in zip(features,image_paths,person_ids.numpy()):
                self.image_features[image_path] = feature
                self.image_person_ids[image_path] = person_id
        with open('./image_features.pkl','wb') as f:
            pickle.dump(self.image_features,   f, protocol=pickle.HIGHEST_PROTOCOL)
        with open('./image_person_ids.pkl','wb') as f:
            pickle.dump(self.image_person_ids, f, protocol=pickle.HIGHEST_PROTOCOL)
    